'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 4_study_runable_parallel.py
* @Time: 2025/7/13
* @All Rights Reserve By Brtc
'''
import dotenv
from langchain_community.chat_models import ChatOpenAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableParallel

# 1、通过  runnable parallel 实现两条链并行 运行
dotenv.load_dotenv()
# 2、定义两个prompt
joke_prompt = ChatPromptTemplate.from_template("请讲一个关于{subject}的冷笑话")
poem_prompt = ChatPromptTemplate.from_template("请写一首关于{subject}的诗")
#3、创建大语言模型
llm = ChatOpenAI(model="gpt-3.5-turbo-16k")
#4、输出解析器， 这个换成  JsonOutputparaser，定定义一个
# Class Jocker(BaseModel):
#     joker_name : str = Field(description = "回答用户冷笑话")
#     punch_line: str = Field(description="这个冷笑话的笑点")

parser = StrOutputParser()
# 5、编排链
joke_chain = joke_prompt|llm|parser
poem_chain = poem_prompt|llm|parser

# 并行链
#RunnableParallel构建了一个，可并行的组件， 并行的 内容是两条链， 一条是poem链 一条是 joke 链，当然可以 跟更多链。
#这个不限制链的数量
#这输出 json  多条链 你们课后试
# 适用于  我 多节点  运行且前后不影响的情况下，可以用 parallel
map_chain = RunnableParallel(joke = joke_chain, poem = poem_chain)
res = map_chain.invoke({"subject": "程序员"})
print(res)

